package main

import "fmt"

/*
给你一个二维数组，每个值表示这块地形的高度，
求该地形能装多少单位的水

预处理数组 + 降维int->bit + 栈 O(m*n * m*n)
or 小根堆 O(m*n log(k))

堆有点像加了一个维度的栈，通过数值大小来控制进出的顺序（而不只是后进先出）
*/

type Stack struct {
	val [][2]int
	top int
}

func main() {
	var arr [][]int = make([][]int, 3)
	arr[0] = append(arr[0], 0, 4, 6, 0)
	arr[1] = append(arr[1], 5, 2, 1, 6)
	arr[2] = append(arr[2], 0, 6, 6, 0)
	//arr[0] = append(arr[0], 3, 3, 3, 3)
	//arr[1] = append(arr[1], 3, 2, 1, 3)
	//arr[2] = append(arr[2], 3, 3, 3, 3)
	fmt.Println(arr)
	fmt.Println(getSum(arr))
}

func getSum(arr [][]int) int {
	bitArr := BToT(arr)
	fmt.Println(bitArr)
	var sum int
	for _, v := range bitArr {
		r := getWaterCount(v)
		sum = sum + r
		fmt.Println(r)
	}
	return sum
}

func getWaterCount(arr [][]int) int {
	m := len(arr)
	n := len(arr[0])
	var count int
	stack := &Stack{val: [][2]int{}, top: -1}

	for i := 1; i < m-1; i++ {
		for j := 1; j < n-1; j++ {

			thisCount := 0
			if arr[i][j] == 0 {
				//压栈
				stack.val = append(stack.val, [2]int{i, j})
				stack.top++
				for stack.top > -1 {
					//弹栈
					cur := stack.val[stack.top]
					stack.top--
					y, x := cur[0], cur[1]
					if arr[y][x] == 0 {
						arr[y][x] = -1  //置为-1标识该轮算过
						if y == 0 || x == 0 || y == m-1 || x == n-1 {
							thisCount = 0  //触达边界，归零
							break
						}
						thisCount++  //弹栈时若可以装水，加一
						upY, upX := cur[0]-1, cur[1]
						downY, downX := cur[0]+1, cur[1]
						leftY, leftX := cur[0], cur[1]-1
						rightY, rightX := cur[0], cur[1]+1

						//压栈
						stack.val = append(stack.val, [2]int{upY, upX})
						stack.val = append(stack.val, [2]int{downY, downX})
						stack.val = append(stack.val, [2]int{leftY, leftX})
						stack.val = append(stack.val, [2]int{rightY, rightX})
						stack.top = stack.top + 4
					}

				}
				count = count + thisCount
			}
			//清空为-1的记录
			for i := 1; i < m-1; i++ {
				for j := 1; j < n-1; j++ {
					if arr[i][j] == -1 {
						arr[i][j] = 0
					}
				}
			}
			//stack = &Stack{val: [][2]int{}, top: -1}  //清空栈

		}
	}

	return count
}

func BToT(arr [][]int) [][][]int {
	var high int
	for _, v := range arr {
		for _, q := range v {
			high = Max(high, q)
		}
	}
	m := len(arr)
	n := len(arr[0])
	var bitArr [][][]int
	for i := 0; i < high; i++ {
		var arr2 [][]int
		for j := 0; j < m; j++ {
			arr2 = append(arr2, make([]int, n))
		}
		bitArr = append(bitArr, arr2)
	}

	for i := 0; i < m; i++ {
		for j := 0; j < n; j++ {
			for h := 0; h < arr[i][j]; h++ {
				bitArr[h][i][j] = 1
			}
		}
	}
	return bitArr
}

func Max(a int, b int) int {
	if a > b {
		return a
	}
	return b
}